package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
)

// AES解密
func AesDecrypt(rawData, key, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// 加密模式
	blockMode := cipher.NewCBCDecrypter(block, iv)

	plainText := make([]byte, len(rawData))
	blockMode.CryptBlocks(plainText, rawData)
	plainText = unPaddingPKCS7(plainText)

	return plainText, nil
}

// AES加密
func AesEncrypt(rawData, key, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// 填充原文
	rawData = paddingPKCS7(rawData, block.BlockSize())
	// 加密模式
	blockMode := cipher.NewCBCEncrypter(block, iv)

	cipherText := make([]byte, len(rawData))
	blockMode.CryptBlocks(cipherText, rawData)

	return cipherText, nil
}

// base64(aes)
func Base64AesEncrypt(rawData, key, iv []byte) (string, error) {
	text, err := AesEncrypt(rawData, key, iv)
	if err != nil {
		return "", err
	}

	return base64.StdEncoding.EncodeToString(text), nil
}

// base64(aes)
func Base64AesDecrypt(rawData string, key, iv []byte) (string, error) {
	cipherText, err := base64.StdEncoding.DecodeString(rawData)
	if err != nil {
		return "", err
	}

	text, err := AesDecrypt(cipherText, key, iv)
	if err != nil {
		return "", err
	}

	return string(text), nil
}

// 补码
func paddingPKCS7(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// 去除补码
func unPaddingPKCS7(data []byte) []byte {
	length := len(data)
	unPadding := int(data[length-1])
	return data[:length-unPadding]
}
